package io.mycat.route.function;

import io.mycat.config.model.rule.RuleAlgorithm;
import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.InputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;

/**
 * Created by LiuChang on 2018/3/27.
 */
public class PartitionByAccount extends AbstractPartitionAlgorithm implements RuleAlgorithm {

    private static final Logger LOGGER = LoggerFactory.getLogger("PartitionByAccount");

    private Map<String, Integer> accountNumMap = new HashMap<>();
    private String xmlFile;
    private Integer defaultNode = 0;

    public void setXmlFile(String xmlFile) {
        this.xmlFile = xmlFile;
    }

    public void setDefaultNode(Integer defaultNode) {
        this.defaultNode = defaultNode;
    }

    @Override
    public void init() {
        try {
            //读取xml配置文件
            InputStream fin = this.getClass().getClassLoader()
                    .getResourceAsStream(xmlFile);
            if (fin == null) {
                throw new RuntimeException("can't find class resource file " + xmlFile);
            }
            Document document = new SAXReader().read(fin);
            Element mysql = document.getRootElement();
            String host = mysql.element("host").getText();
            String port = mysql.element("port").getText();
            String database = mysql.element("database").getText();
            String username = mysql.element("username").getText();
            String password = mysql.element("password").getText();
            String tableName = mysql.element("tableName").getText();
            String accountColumn = mysql.element("accountColumn").getText();
            String numColumn = mysql.element("numColumn").getText();
            LOGGER.info("读取配置文件：host(" + host + "),port(" + port + "),database(" + database + "),username(" + username + "),password(" + password + "),tablename(" + tableName + ")");
            //    读取表中的数据
            String url = "jdbc:mysql://" + host + ":" + port + "/" + database;
            String sql = "SELECT " + accountColumn + "," + numColumn + " FROM " + tableName;
            Connection connection = DriverManager.getConnection(url, username, password);
            Statement statement = connection.createStatement();
            ResultSet resultSet = statement.executeQuery(sql);
            if (resultSet != null) {
                while (resultSet.next()) {
                    accountNumMap.put(resultSet.getString(accountColumn), resultSet.getInt(numColumn));
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    @Override
    public Integer calculate(String columnValue) {
        Integer partition = accountNumMap.get(columnValue);
        if (partition == null) {
            return defaultNode;
        }
        return partition;
    }
}
